"""Regression model"""

import numpy as np
import matplotlib.pyplot as plt
from giw import *
from scipy.constants import pi


class Regmodel(Giw):
    """
    Basic Bayesian (auto)regression model.
    
    There will be yet models with various forgetting methods
    implemented, which will inherit from this model.
    

    Attributes
    ----------
    _log : dict
        Logs
    data_real : array
        Array of real data for modelling
    mixtools_compat : bool
        Compatibility mode for likelihoods
    delta_log_likelihood : float
        Increment of log-likelihood
    regression_delay : int
        Delayed regression


    Methods
    -------
    __init__(mat_v=None, nu=1, desc=None, mixtools_compat=False)
        constructor
    plot_est_theta_hist()
        plots estimated regression coefficients
    model_update(data)
        model_update model pdf
    regress(regression_vec, y_true=None, delay=0)
        regression using regression vector

    Properties
    ----------
    model_log_likelihood
        log_likelihood of the model, either absolute or mixtools-like
    est_theta_hist
        history of regr. coefficients estimations
    log
        get the log dictionary
    """

    def __init__(self, mat_v=None, nu=1, desc=None, mixtools_compat=False):
        """
        Constructor

        Parameters
        ----------
        mat_v : np.array
            Information matrix
        nu : float
            Degrees of freedom
        desc : str
            Description
        mixtools_compat : bool
            Compatibility with mixtools' loglik
        """
        super(Regmodel, self).__init__(mat_v, nu, desc)
        self._model_log_likelihood = 0
        self._log = {'est_theta_hist': [],
                     'rgr_hist': [],
                     'y_true': [],
                     'y_residual': [],
                     'y_predicted': []}
        self.mixtools_compat = mixtools_compat
        self.delta_log_likelihood = 0
        self.regression_delay = 0

    def model_update(self, data):
        """
        Data update of model
        """
        prior_log_likelihood = self.log_likelihood
        self.update(data)
        self._log['est_theta_hist'].append(self.est_theta)
        posterior_log_likelihood = self.log_likelihood

        self.delta_log_likelihood = posterior_log_likelihood - prior_log_likelihood
        self._model_log_likelihood += self.delta_log_likelihood

    @property
    def log(self):
        return self._log

    @property
    def est_theta_hist(self):
        return self._log['est_theta_hist']

    @property
    def model_log_likelihood(self):
        return self._model_log_likelihood

    def plot_est_theta_hist(self):
        fig = plt.figure(1)
        ax = fig.add_subplot(111)
        ax.grid(True)
        ax.hold(True)
        ax.plot(np.array(self._log['est_theta_hist'])[:,0], '--', color='black')
        ax.plot(np.array(self._log['est_theta_hist'])[:,1], '-', color='black')
        ax.set_title("Estimated regression coefficients")

    def regress2(self, regression_vec):
        #print self.est_theta, ' and ', regression_vec, ' =>', np.dot(self.est_theta, regression_vec)
        return np.dot(self.est_theta, regression_vec)

    def regress(self, regression_vec, y_true=None, delay=0):
        """
        Model regression from regression vector. 
        
        Parameters
        ----------
        regression_vec : np.array
            Regression vector
        y_true : float
            True data. If supplied, the residue is returned as well.
        delay : int
            Traffic delay, usually unused. Only first setting is applied!!!
        """
        if delay != 0:
            self.regression_delay = delay

            try:
                self.stack
            except AttributeError:
                from Queue import Queue
                self.stack = Queue(delay + 1)

            self.stack.put_nowait(self.est_theta)

            if self.stack.full():
                regression_coefs = self.stack.get()
            else:
            	regression_coefs = np.empty(regression_vec.shape)
            	regression_coefs[0] = np.nan
        else:
        	regression_coefs = self.est_theta

        y = np.dot(regression_vec, regression_coefs)
        self._log['y_predicted'].append(y)
        self._log['rgr_hist'].append(regression_vec)
        
        if y_true != None:
        	self._log['y_true'].append(y_true)
        	self._log['y_residual'].append(y_true - y)


#--------------------------------------------

if __name__ == "__main__":
    #Initial conditions
    regression_coefs = np.array([0.5, -0.9])
    data_true = [.4]

    #Generate AR(1) data
    for i in xrange(10):
        data_true.append(np.dot([data_true[i], 1], regression_coefs) + np.random.normal(scale=0.01)) 
    print data_true

    #Prepare data for modelling
    datamatrix = data_matrix(data_true, 1)

    #Prepare regressive model
    V = np.eye(3) * 0.01
    V[0,0] = 0.1
    nu = 10
    model = Regmodel(V, nu, mixtools_compat=True)

    #Modelling - regression starts from noninformative prior
    loglikelihood_prior = model.model_log_likelihood

    check = []
    for dt in datamatrix:
        regressor = dt[1:]
        model.regress(regressor, dt[0], delay=1)
        check.append(model.regress2(regressor))
        model.model_update(dt)

    log = model.log
    print "Regression coefs.\n", log['rgr_hist']
    print "Residues\n", log['y_residual']

    loglikelihood_posterior = model.model_log_likelihood

    #print "Information matrix\n", model.v
    #print "Degrees of freedom: ", model.nu

    print "Loglik. cumulative", model._model_log_likelihood
    print "Loglik. post - prior", loglikelihood_posterior - loglikelihood_prior

    model.plot_est_theta_hist()
    fig = plt.figure(2)
    ax = fig.add_subplot(211)
    ax.set_title('Residues')
    ax.plot(log['y_residual'])
    ax = fig.add_subplot(212)
    ax.set_title('Predictions')
    ax.plot(log['y_true'])
    ax.plot(log['y_predicted'], 'oy')
    ax.plot(check, 'xr', markersize=10)
    plt.show()
